package com.zipkin.demo.config;

import brave.CurrentSpanCustomizer;
import brave.SpanCustomizer;
import brave.Tracing;
import brave.baggage.BaggageField;
import brave.baggage.BaggagePropagation;
import brave.baggage.BaggagePropagationConfig;
import brave.baggage.CorrelationScopeConfig;
import brave.context.slf4j.MDCScopeDecorator;
import brave.http.HttpTracing;
import brave.okhttp3.TracingInterceptor;
import brave.propagation.B3Propagation;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.ThreadLocalCurrentTraceContext;
import brave.servlet.TracingFilter;
import brave.spring.webmvc.SpanCustomizingAsyncHandlerInterceptor;
import okhttp3.OkHttpClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.web.client.RestTemplateCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.http.client.OkHttp3ClientHttpRequestFactory;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter;
import zipkin2.reporter.Sender;
import zipkin2.reporter.brave.AsyncZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import javax.servlet.Filter;

/**
 * 针对mvc controller 和 restTemplate 的 zipkin客户端配置
 */
@Configuration
@Import(SpanCustomizingAsyncHandlerInterceptor.class)
public class ZipkinClientConfiguration implements WebMvcConfigurer {

    public static final BaggageField USER_NAME = BaggageField.create("userName");

    /**
     * Configuration for how to send spans to Zipkin
     */
    @Bean
    Sender sender(
            @Value("${brave.zipkin.baseUrl:http://127.0.0.1:9411}/api/v2/spans") String zipkinEndpoint) {
        return OkHttpSender.create(zipkinEndpoint);
    }


    /**
     * Configuration for how to buffer spans into messages for Zipkin
     *
     * @param sender
     * @return
     */
    @Bean
    AsyncZipkinSpanHandler zipkinSpanHandler(Sender sender) {
        return AsyncZipkinSpanHandler.create(sender);
    }

    /**
     * 配置跟踪过程中的Trace信息
     * Controls aspects of tracing such as the service name that shows up in the UI
     */
    @Bean
    Tracing tracing(
            @Value("${brave.localServiceName:${spring.application.name}}") String serviceName,
            @Value("${brave.supportsJoin:true}") boolean supportsJoin,
            @Value("${brave.traceId128Bit:false}") boolean traceId128Bit,
            Propagation.Factory propagationFactory,
            CurrentTraceContext currentTraceContext,
            AsyncZipkinSpanHandler zipkinSpanHandler) {
        return Tracing.newBuilder()
                .localServiceName(serviceName)
                .supportsJoin(supportsJoin)
                .traceId128Bit(traceId128Bit)
                .propagationFactory(propagationFactory)
                .currentTraceContext(currentTraceContext)
                .addSpanHandler(zipkinSpanHandler).build();
    }


    @Bean
    CurrentTraceContext.ScopeDecorator correlationScopeDecorator() {
        return MDCScopeDecorator.newBuilder()
                .add(CorrelationScopeConfig.SingleCorrelationField.create(USER_NAME)).build();
    }


    /**
     * Propagates trace context between threads.
     */
    @Bean
    CurrentTraceContext currentTraceContext(CurrentTraceContext.ScopeDecorator scopeDecorator) {
        return ThreadLocalCurrentTraceContext.newBuilder()
                .addScopeDecorator(scopeDecorator)
                .build();
    }


    /**
     * Configures propagation for {@link #USER_NAME}, using the remote header "user_name"
     *
     * @return
     */
    @Bean
    Propagation.Factory propagationFactory() {
        return BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
                .add(BaggagePropagationConfig.SingleBaggageField.newBuilder(USER_NAME).addKeyName("user_name").build())
                .build();
    }

    /**
     * 注入可定制的Span
     */
    @Bean
    SpanCustomizer spanCustomizer(Tracing tracing) {
        return CurrentSpanCustomizer.create(tracing);
    }

    /**
     * 决定如何命名和标记span。 默认情况下，它们的名称与http方法相同
     */
    @Bean
    HttpTracing httpTracing(Tracing tracing) {
        return HttpTracing.create(tracing);
    }

    /**
     * 导入过滤器，该过滤器中会为http请求创建span
     */
    @Bean
    Filter tracingFilter(HttpTracing httpTracing) {
        return TracingFilter.create(httpTracing);
    }


    @Bean
    RestTemplateCustomizer useOkHttpClient(final OkHttpClient okHttpClient) {
        return new RestTemplateCustomizer() {
            @Override
            public void customize(RestTemplate restTemplate) {
                restTemplate.setRequestFactory(new OkHttp3ClientHttpRequestFactory(okHttpClient));
            }
        };
    }

    /**
     * Creates client spans for HTTP requests.
     *
     * @param httpTracing
     * @return
     */
    @Bean
    OkHttpClient tracedOkHttpClient(HttpTracing httpTracing) {
        return new OkHttpClient.Builder()
                .addNetworkInterceptor(TracingInterceptor.create(httpTracing))
                .build();
    }


    /**
     * Adds MVC Controller tags to server spans
     */
    @Bean
    WebMvcConfigurer tracingWebMvcConfigurer(
            final SpanCustomizingAsyncHandlerInterceptor webMvcTracingCustomizer) {
        return new WebMvcConfigurerAdapter() {
            /** Adds application-defined web controller details to HTTP server spans */
            @Override
            public void addInterceptors(InterceptorRegistry registry) {
                registry.addInterceptor(webMvcTracingCustomizer);
            }
        };
    }


}
